from flask import request, Blueprint,render_template, redirect, url_for
from datetime import datetime
from models import MonthlyBudget, Expense, db
from decimal import Decimal

budget_bp = Blueprint('budget', __name__, url_prefix='/budget')

@budget_bp.route("/", methods=["GET", "POST"])
def set_budget():
    if request.method == "POST":
        try:
            budget_amount = Decimal(request.form["budget"]).quantize(Decimal('0.00'))
            month = int(request.form.get("month", datetime.now().month))
            year = int(request.form.get("year", datetime.now().year))

            # Check if budget already exists for this month/year
            existing_budget = MonthlyBudget.query.filter_by(
                month=month, year=year
            ).first()

            if existing_budget:
                # 更新剩余预算为新预算与已使用预算的差值
                existing_budget.budget = Decimal(str(existing_budget.budget)) + budget_amount
                existing_budget.remaining = Decimal(str(existing_budget.remaining)) + budget_amount
            else:
                new_budget = MonthlyBudget(
                    month=month,
                    year=year,
                    budget=budget_amount,
                    remaining=budget_amount,
                )
                db.session.add(new_budget)

            db.session.commit()
            return redirect(url_for("budget.set_budget"))
        except Exception as e:
            db.session.rollback()
            return f"设置预算时出错: {str(e)}", 500

    # Get current month and year for default values
    now = datetime.now()
    current_month = now.month
    current_year = now.year

    # Get all budgets for current month/year
    current_budgets = MonthlyBudget.query.filter_by(
        month=current_month, year=current_year
    ).all()

    # Get available months and years for history view
    budget_months = (
        db.session.query(MonthlyBudget.month, MonthlyBudget.year)
        .distinct()
        .order_by(MonthlyBudget.year.desc(), MonthlyBudget.month.desc())
        .all()
    )

    expense_months = (
        db.session.query(Expense.month, Expense.year)
        .distinct()
        .order_by(Expense.year.desc(), Expense.month.desc())
        .all()
    )

    # Combine and deduplicate months/years
    all_months = list(set(budget_months + expense_months))
    all_months.sort(key=lambda x: (x[1], x[0]), reverse=True)

    return render_template(
        "budget.html",
        current_month=current_month,
        current_year=current_year,
        current_budgets=current_budgets,
        all_months=all_months,
    )

@budget_bp.route("/history")
def budget_history():
    # 获取查询参数
    month = request.args.get("month", type=int)
    year = request.args.get("year", type=int)

    # 如果没有提供月份和年份，使用当前月份
    if not month or not year:
        now = datetime.now()
        month = now.month
        year = now.year

    # 获取指定月份的预算和支出
    budgets = MonthlyBudget.query.filter_by(month=month, year=year).all()

    # 获取指定月份的支出总额
    expenses = (
        db.session.query(db.func.sum(Expense.amount).label("total"))
        .filter_by(month=month, year=year)
        .all()
    )

    # 获取可用的月份列表
    budget_months = (
        db.session.query(MonthlyBudget.month, MonthlyBudget.year)
        .distinct()
        .order_by(MonthlyBudget.year.desc(), MonthlyBudget.month.desc())
        .all()
    )

    expense_months = (
        db.session.query(Expense.month, Expense.year)
        .distinct()
        .order_by(Expense.year.desc(), Expense.month.desc())
        .all()
    )

    # 合并并去重
    all_months = list(set(budget_months + expense_months))
    all_months.sort(key=lambda x: (x[1], x[0]), reverse=True)

    return render_template(
        "budget_history.html",
        budgets=budgets,
        expenses=expenses,
        selected_month=month,
        selected_year=year,
        all_months=all_months,
    )

# 用于删除预算
@budget_bp.route("/delete/<int:budget_id>", methods=["POST"])
def delete_budget(budget_id):
    # 尝试执行以下代码
    try:
        # 根据预算id查找预算
        budget = MonthlyBudget.query.get_or_404(budget_id)
        # 从数据库中删除该预算
        db.session.delete(budget)
        # 提交更改
        db.session.commit()
        # 重定向到设置预算页面
        return redirect(url_for("budget.set_budget"))
    # 如果发生异常，执行以下代码
    except Exception as e:
        # 回滚更改
        db.session.rollback()
        # 返回错误信息
        return f"删除预算时出错: {str(e)}", 500

# 添加重新计算路由
@budget_bp.route("/recalculate", methods=["POST"])
def recalculate_budget():
    try:
        # 获取当前月份和年份
        now = datetime.now()
        current_month = now.month
        current_year = now.year
        
        # 获取当月的所有预算记录
        budgets = MonthlyBudget.query.filter_by(
            month=current_month, year=current_year
        ).all()
        
        # 获取当月的所有支出总额
        total_expense_result = db.session.query(db.func.sum(Expense.amount)).filter_by(
            month=current_month, year=current_year
        ).scalar()
        
        total_expense = Decimal(str(total_expense_result)) if total_expense_result else Decimal('0.00')
        total_expense = total_expense.quantize(Decimal('0.00'))
        
        # 重新计算每个预算的剩余金额
        for budget in budgets:
            # 新的剩余金额 = 总预算 - 总支出
            new_remaining = (Decimal(str(budget.budget)) - total_expense).quantize(Decimal('0.00'))
            budget.remaining = new_remaining
            
        # 提交更改
        db.session.commit()
        
        # 重定向到预算页面
        return redirect(url_for("budget.set_budget"))
    except Exception as e:
        # 回滚更改
        db.session.rollback()
        # 返回错误信息
        return f"重新计算预算时出错: {str(e)}", 500